///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file energy_accumulator.h
///
/// \brief   header file for the energy accumulator class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef PCPS_ENERGY_ACCUMULATOR_HEADER
#define PCPS_ENERGY_ACCUMULATOR_HEADER

#include<src/pcps.h>
#include<src/input.h>
#include<src/configuration.h>
#include<src/accumulator.h>
#include<src/coeff_obj.h>

namespace pcps {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  a class to accumulate energies during a markov chain evolution
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////
  template <class S> class EnergyAccumulator : public pcps::Accumulator<S> {

    private:

      /// \brief  the number of alpha sites
      int _nas;

      /// \brief  the number of alpha particles
      int _nap;

      /// \brief  the number of unoccupied alpha sites
      int _nau;

      /// \brief  the maximum lenth for which to compute autocorrelations
      int _max_ac_length;

      /// \brief  the number of cholesky vectors
      int _nch;

      /// \brief  pointer to the coefficient object to use
      boost::shared_ptr<pcps::CoeffObj<S> > _co;

      /// \brief  pointer to the Hamiltonian
      const pcps::Hamiltonian<S> * _ham;

      /// \brief  whether we are sampling w.r.t. to the wavefunction or the local rdm
      bool _wfn_sampling;

      /// \brief  whether to weight the local rdm elements by the [pq,pq] integrals
      bool _pqpq_guiding;

      /// \brief  whether we are to print detailed information about each sample
      bool _print;

      /// \brief  energy
      double _energy;

      /// \brief  energy
      double _pqqp_multiplier;

      /// \brief  stores one electron energy for possible printing
      double _oei_contribution;

      /// \brief  stores pqqp energy for possible printing
      double _pqqp_energy;

      /// \brief  stores pppp energy for possible printing
      double _pppp_energy;

      /// \brief  work space to hold ratios between wavefunction values for single alpha excitations
      std::vector<S> _a_ratios;

      /// \brief  work space to hold ratios between wavefunction values for single beta excitations
      std::vector<S> _b_ratios;

      /// \brief  vector to hold the contraction of the cholesky vectors with the alpha ratio matrix
      std::vector<S> _a_contraction;

      /// \brief  vector to hold the contraction of the cholesky vectors with the beta ratio matrix
      std::vector<S> _b_contraction;

      /// \brief  vector to hold the contraction of the cholesky vectors with the diagonal of the alpha ratio matrix
      std::vector<S> _a_diag_cont;

      /// \brief  vector to hold the contraction of the cholesky vectors with the diagonal of the beta ratio matrix
      std::vector<S> _b_diag_cont;

      /// \brief  a vector to hold the slice of the one electron integrals needed for contractions with ratio matrices
      std::vector<S> _oei_slice;

      /// \brief  a vector to hold the slice of the cholesky vector matrix needed for contractions with ratio matrices
      std::vector<S> _chol_slice;

      /// \brief  a vector to hold the slice of the one electron integrals needed for contractions with alpha ratio matrices
      std::vector<S> _a_oei_slice;

      /// \brief  a vector to hold the slice of the cholesky vector matrix needed for contractions with alpha ratio matrices
      std::vector<S> _a_chol_slice;

      /// \brief  a vector to hold the slice of the one electron integrals needed for contractions with alpha ratio matrices
      std::vector<S> _b_oei_slice;

      /// \brief  a vector to hold the slice of the cholesky vector matrix needed for contractions with alpha ratio matrices
      std::vector<S> _b_chol_slice;

      /// \brief  the [pq,pq] two electron integrals
      std::vector<S> _pqpq_tei;

      /// \brief  a history of the sampled local energies
      std::vector<double> _history;

      /// \brief  the configuration we currently have our slices for
      pcps::Configuration<S> _config;

      void get_ham_slice(const pcps::Input & userinp, const pcps::Configuration<S> & config, const int b);

      void update_slices(const pcps::Configuration<S> & config);

      void init_slices();

      void calculate_a(const pcps::Input & userinp,
                       const double w,
                       const pcps::Configuration<S> & config,
                       const int b,
                       const std::vector<S> & ratios,
                       std::vector<S> & contraction,
                       std::vector<S> & diag_cont,
                       const std::vector<S> & opp_spin_ratios);

      /// \brief  binary operator function object used in computing the variance
      struct VarBinOp {

        double m;

        VarBinOp(double mean) : m(mean) {}
      
        double operator()(double x, double y) { return x + y*y - m*m; }
        
      };

      /// \brief  binary operator function object used in computing the autocorrelation
      struct AutoCorrBinOp {

        double m;
        int t;

        AutoCorrBinOp(double mean, int offset_time) : m(mean), t(offset_time) {}

        double operator()(const double x, const double & y) {
          return x + ( y - m ) * ( *(&y-t) - m );
        }

      };

    public:

      EnergyAccumulator(const pcps::Input & userinp,
                        const pcps::Hamiltonian<S> & ham,
                        boost::shared_ptr<pcps::CoeffObj<S> > & co);

      /// \brief  return the energy
      double energy() const { return _energy; }

      /// \brief  return the sampling type flag
      bool wfn_sampling() const { return _wfn_sampling; }

      /// \brief  return the sampling type flag by reference
      bool & wfn_sampling() { return _wfn_sampling; }

      /// \brief  return the print flag
      bool print() const { return _print; }

      /// \brief  return the print flag by reference
      bool & print() { return _print; }

      /// \brief  return the pqqp multiplier
      double pqqp_multiplier() const { return _pqqp_multiplier; }

      /// \brief  return the pqqp multiplier by reference
      double & pqqp_multiplier() { return _pqqp_multiplier; }

      void print_header();
      void print_statistics();

      void child_accumulate(const pcps::Input & userinp, const double w, const pcps::Configuration<S> & config);

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief finalizes the accumulation (this may include for example communicating the results)
      ///
      /// \param[in]      userinp  the user's input options
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_finalize(const pcps::Input & userinp) {
        double x = 0.0;
        //pcps::mpi::allreduce(&_norm_ratio, &x, 1, MPI::SUM);
        //_norm_ratio = x / this->_total_weight;
        //x = 0.0;
        x = 0.0;
        pcps::mpi::allreduce(&_energy, &x, 1, MPI::SUM);
        _energy = x / this->_total_weight;
        if (_wfn_sampling) {
          x = 0.0;
          pcps::mpi::allreduce(&_pqqp_multiplier, &x, 1, MPI::SUM);
          _pqqp_multiplier = x / this->_total_weight;
        }
        //std::cout << "total weight = " << this->_total_weight << std::endl;
        //_energy = _ham->e_cor_nuc() + x / this->_total_weight / _norm_ratio;
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  resets the accumulator
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void child_reset(const pcps::Configuration<S> & config) {
        _energy = 0.0;
        if (_wfn_sampling)
          _pqqp_multiplier = 0.0;
        else if (!_pqpq_guiding)
          _pqqp_multiplier = 1.0;
        _history.clear();
        _config = config;
        this->init_slices();
        //_norm_ratio = 0.0;
      }

  };

//  ///////////////////////////////////////////////////////////////////////////////////////////////////
//  /// \brief  the base class for accumulator objects that collect data during a monte carlo sample
//  ///
//  ///////////////////////////////////////////////////////////////////////////////////////////////////
//  template <class T> class EnergyAccumulator : public pcps::Accumulator<T> {
//
//    private:
//
//      /// \brief  pointer to the coefficient object used as the guiding and value function
//      boost::shared_ptr<pcps::CoeffObj<T> > _co;
//
//      /// \brief  pointer to the user's input options
//      const pcps::Input * _userinp;
//
//      /// \brief  pointer to the Hamiltonian
//      const pcps::Hamiltonian<T> * _ham;
//
//      /// \brief  energy
//      double _energy;
//
//      /// \brief  ratio between the square norms of the wavefunction without and with noise
//      double _norm_ratio;
//
//      /// \brief  work space to hold ratios between wavefunction values for single alpha excitations
//      std::vector<T> _a_ratios;
//
//      /// \brief  work space to hold ratios between wavefunction values for single beta excitations
//      std::vector<T> _b_ratios;
//
//      /// \brief  vector to hold the contraction of the cholesky vectors with the alpha ratio matrix
//      std::vector<T> _a_contraction;
//
//      /// \brief  vector to hold the contraction of the cholesky vectors with the beta ratio matrix
//      std::vector<T> _b_contraction;
//
//      /// \brief  for holding the local configuration's value of <n| H_1 |Psi>
//      double _oe_local;
//
//      /// \brief  for holding the local configuration's value of sum_pq <Psi| p+ q |n><n| q+ p |Psi>
//      double _local_2rdm;
//
//      std::vector<T> _a_diag_contraction;
//      std::vector<T> _b_diag_contraction;
//      std::vector<T> _a_offd_contraction;
//      std::vector<T> _b_offd_contraction;
//
//      void calculate_a(const pcps::Input & userinp,
//                       const double w,
//                       const pcps::Configuration<T> & config,
//                       const int b,
//                       const std::vector<T> & ratios,
//                       std::vector<T> & contraction,
//                       std::vector<T> & diag_cont,
//                       std::vector<T> & offd_cont);
//
//    public:
//
//      /// \brief  constructor
//      EnergyAccumulator(const pcps::Input & userinp,
//                        const pcps::Hamiltonian<T> & ham,
//                        boost::shared_ptr<pcps::CoeffObj<T> > & co)
//        : Accumulator<T>(),
//          _co(co),
//          _userinp(&userinp),
//          _ham(&ham),
//          _energy(0.0),
//          _norm_ratio(0.0),
//          _a_ratios( (userinp.nsites()/2) * (userinp.nsites()/2), pcps::zero<T>() ),
//          _b_ratios( (userinp.nsites()/2) * (userinp.nsites()/2), pcps::zero<T>() ),
//          _a_contraction(ham.n_cholesky(), pcps::zero<T>()),
//          _b_contraction(ham.n_cholesky(), pcps::zero<T>()),
//          _a_diag_contraction(ham.n_cholesky(), pcps::zero<T>()),
//          _b_diag_contraction(ham.n_cholesky(), pcps::zero<T>()),
//          _a_offd_contraction(ham.n_cholesky(), pcps::zero<T>()),
//          _b_offd_contraction(ham.n_cholesky(), pcps::zero<T>())
//      {}
//
//      /// \brief  return the energy
//      double energy() const { return _energy; }
//
//      void child_accumulate(const pcps::Input & userinp, const double w, const pcps::Configuration<T> & config);
//
//      ///////////////////////////////////////////////////////////////////////////////////////////////////
//      /// \brief finalizes the accumulation (this may include for example communicating the results)
//      ///
//      /// \param[in]      userinp  the user's input options
//      ///
//      ///////////////////////////////////////////////////////////////////////////////////////////////////
//      void child_finalize(const pcps::Input & userinp) {
//        double x = 0.0;
//        pcps::mpi::allreduce(&_norm_ratio, &x, 1, MPI::SUM);
//        _norm_ratio = x / this->_total_weight;
//        x = 0.0;
//        pcps::mpi::allreduce(&_energy, &x, 1, MPI::SUM);
//        _energy = _ham->e_cor_nuc() + x / this->_total_weight / _norm_ratio;
//      }
//
//      ///////////////////////////////////////////////////////////////////////////////////////////////////
//      /// \brief  resets the accumulator
//      ///
//      ///////////////////////////////////////////////////////////////////////////////////////////////////
//      void child_reset() {
//        _energy = 0.0;
//        _norm_ratio = 0.0;
//      }
//
//  };

}

#endif
